#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : lei.wang
import anyio
import fastapi
import alembic
import asyncio
import redis.asyncio as redis
from datetime import datetime
from pathlib import Path
from collections.abc import AsyncGenerator, Callable
from contextlib import asynccontextmanager
from typing import Any, Union
from alembic.config import Config
from fastapi import APIRouter, Depends, FastAPI, Request
from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html
from fastapi.openapi.utils import get_openapi
from fastapi.templating import Jinja2Templates
from fastapi.middleware.cors import CORSMiddleware
from app.core.db.database import Base, async_engine as engine
from app.models import *
from app.core.config import config
from app.core.pkg.log import logger

# 允许的来源
allowed_origins = ["*"]  # 允许所有来源，实际环境中应该指定具体域名

# 导入系统环境变量中设置的允许来源
import os
# 从环境变量获取允许来源
if os.environ.get("ALLOW_ORIGINS"):
    allowed_origins = os.environ.get("ALLOW_ORIGINS").split(",")

templates = Jinja2Templates(directory="templates")

# -------------- local storage --------------
CURRENT_DIR = Path(__file__).absolute().parent.parent
CURRENT_DATE_STR = datetime.now().strftime('%Y%m%d')
LOCAL_UPLOAD_DIR = CURRENT_DIR / 'uploads' / CURRENT_DATE_STR
VIDEO_UPLOAD_DIR = CURRENT_DIR / 'uploads' / 'video' / CURRENT_DATE_STR
AUDIO_UPLOAD_DIR = CURRENT_DIR / 'uploads' / 'audio' / CURRENT_DATE_STR


async def _init_local_storage() -> None:
    if not LOCAL_UPLOAD_DIR.exists():
        LOCAL_UPLOAD_DIR.mkdir(parents=True)
        VIDEO_UPLOAD_DIR.mkdir(parents=True)
        AUDIO_UPLOAD_DIR.mkdir(parents=True)


# -------------- local storage --------------


# -------------- database --------------
async def create_tables() -> None:
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


# -------------- application --------------
async def set_threadpool_tokens(number_of_tokens: int = 100) -> None:
    limiter = anyio.to_thread.current_default_thread_limiter()
    limiter.total_tokens = number_of_tokens


# 封装 Alembic 迁移为异步函数
async def run_alembic_migrations():
    path = CURRENT_DIR / 'alembic.ini'
    loop = asyncio.get_running_loop()
    # 在新线程中执行同步的 Alembic 迁移操作
    await loop.run_in_executor(None, lambda: alembic.config.command.upgrade(Config(path), 'head'))


def lifespan_factory(create_tables_on_start: bool = True):
    """Factory to create a lifespan async context manager for a FastAPI app."""

    @asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncGenerator:
        await set_threadpool_tokens()

        await _init_local_storage()
        logger.success('Initializing local storage success!')

        if create_tables_on_start:
            await create_tables()

        # 检查是否启用迁移配置并执行迁移
        # if config.get("ENABLED_MIGRATIONS"):
        #     logger.success('Starting migrations')
        #     try:
        #         await run_alembic_migrations()
        #         logger.success('Finished migrations')
        #     except Exception as e:
        #         logger.error(f"Error running migrations: {e}")

        from app.core.pkg.log import init_logger

        init_logger(CURRENT_DIR / 'logs')

        # 初始化 Redis 连接
        redis_client = None
        if config.get("redis").get("enable"):
            try:
                redis_client = redis.Redis(
                    host=config.get('redis').get('host'),
                    port=config.get('redis').get('port'),
                    db=config.get('redis').get('db'),
                    password=config.get('redis').get('password')
                )
                # 测试连接
                await redis_client.ping()
                app.state.redis = redis_client
                logger.success("Redis connection established.")
            except Exception as e:
                logger.error(f"Failed to connect to Redis: {e}")
        yield

        # 关闭 Redis 连接
        if redis_client:
            await redis_client.close()
            logger.success("Redis connection closed.")

    return lifespan


# -------------- application --------------
def create_application(
        router: APIRouter,
        create_tables_on_start: bool = True,
        **kwargs: Any,
) -> FastAPI:
    # --- before creating application ---
    to_update = {
        "title": config.get("APP_NAME", default=''),
        "description": config.get("APP_DESCRIPTION", default=''),
        "contact": {"name": config.get('CONTACT_NAME', default=''), "email": config.get('CONTACT_EMAIL', default='')},
        "license_info": {"name": config.get('LICENSE_NAME', default=''), "url": config.get('LICENSE_URL', default='')},
    }
    kwargs.update(to_update)

    kwargs.update({"docs_url": None, "redoc_url": None, "openapi_url": None})

    lifespan = lifespan_factory(create_tables_on_start=create_tables_on_start)

    application = FastAPI(lifespan=lifespan, **kwargs)
    
    # 增强的CORS中间件配置
    application.add_middleware(
        CORSMiddleware,
        allow_origins=allowed_origins,  # 从环境变量获取或默认允许所有
        allow_credentials=True,
        allow_methods=["*"],  # 允许所有方法
        allow_headers=["*"],  # 允许所有头
        expose_headers=["Content-Length", "Content-Range", "Authorization"],
        max_age=3600,  # 预检请求的缓存时间
    )

    @application.get('/')
    def index(request: Request):
        return templates.TemplateResponse('index.html', {'request': request})

    application.include_router(router)

    docs_router = APIRouter()

    @docs_router.get("/docs", include_in_schema=False)
    async def get_swagger_documentation() -> fastapi.responses.HTMLResponse:
        return get_swagger_ui_html(openapi_url="/openapi.json", title="docs")

    @docs_router.get("/redoc", include_in_schema=False)
    async def get_redoc_documentation() -> fastapi.responses.HTMLResponse:
        return get_redoc_html(openapi_url="/openapi.json", title="docs")

    @docs_router.get("/openapi.json", include_in_schema=False)
    async def openapi() -> dict[str, Any]:
        out: dict = get_openapi(title=application.title, version=application.version, routes=application.routes)
        return out

    application.include_router(docs_router)

    return application
